package com.lovely602.ai.alibaba.config;


import com.alibaba.cloud.ai.graph.GraphRepresentation;
import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.KeyStrategyFactory;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import com.lovely602.ai.alibaba.entity.Product;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.HashMap;
import java.util.Map;

import static com.alibaba.cloud.ai.graph.StateGraph.END;
import static com.alibaba.cloud.ai.graph.StateGraph.START;
import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;

/**
 *
 * @author lizhixing
 */
@Configuration
public class ProductGraphConfiguration {


    @Bean
    public StateGraph productAnalysisGraph(ChatClient.Builder chatClientBuilder) throws GraphStateException {
        ChatClient client = chatClientBuilder
                .defaultAdvisors(new SimpleLoggerAdvisor())
                .build();

        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> keyStrategyHashMap = new HashMap<>();
            keyStrategyHashMap.put("productDesc", new ReplaceStrategy());
            keyStrategyHashMap.put("slogan", new ReplaceStrategy());
            keyStrategyHashMap.put("productSpec", new ReplaceStrategy());
            keyStrategyHashMap.put("finalProduct", new ReplaceStrategy());
            return keyStrategyHashMap;
        };

        NodeAction marketingCopyNode = state -> {
            String productDesc = (String) state.value("productDesc").orElseThrow();
            String slogan = client.prompt()
                    .user("Generate a catchy slogan for a product with the following description: " + productDesc)
                    .call()
                    .content();
            return Map.of("slogan", slogan);
        };

        NodeAction specificationExtractionNode = state -> {
            String productDesc = (String) state.value("productDesc").orElseThrow();
            Product productSpec = client.prompt()
                    .user("Extract product specifications from the following description: " + productDesc)
                    .call()
                    .entity(Product.class);
            return Map.of("productSpec", productSpec);
        };

        NodeAction mergeNode = state -> {
            String slogan = (String) state.value("slogan").orElseThrow();
            Product productSpec = (Product) state.value("productSpec").orElseThrow();
            Product finalProduct = new Product(slogan, productSpec.material(), productSpec.colors(), productSpec.season());
            return Map.of("finalProduct", finalProduct);
        };

        StateGraph graph = new StateGraph("ProductAnalysisGraph", keyStrategyFactory);
        graph.addNode("marketingCopy", node_async(marketingCopyNode))
                .addNode("specificationExtraction", node_async(specificationExtractionNode))
                .addNode("merge", node_async(mergeNode))
                .addEdge(START, "marketingCopy")
                .addEdge(START, "specificationExtraction")
                .addEdge("marketingCopy", "merge")
                .addEdge("specificationExtraction", "merge")
                .addEdge("merge", END);

        GraphRepresentation representation = graph.getGraph(GraphRepresentation.Type.PLANTUML, "Product Analysis Graph");
        System.out.println("\n=== Product Analysis Graph UML Flow ===");
        System.out.println(representation.content());
        System.out.println("======================================\n");

        return graph;
    }

}
